#!/usr/bin/env python

from collections import Counter

class HuffmanTreeNode:
    def __init__(self, item = None, freq = 0, lchild = None, rchild = None):
        self.item = item
        self.freq = freq
        self.lchild = lchild
        self.rchild = rchild

    def is_leaf(self):
        return self.lchild is None and self.rchild is None

class HuffmanTree:
    def __init__(self):
        self._initialize()
    
    def _initialize(self):
        self.codebook = {}
        self.text = None
        self._coding_tree = None

    def update(self, text):
        self._initialize()
        self.text = text
        self._build_huffman_tree()

    def _build_huffman_tree(self):
        codes = Counter(self.text)

        nodes = [HuffmanTreeNode(key, value) for key, value in codes.items()]
        while len(nodes) > 1:
            nodes.sort(key = lambda x : x.freq, reverse=True)
            n1 = nodes.pop()
            n2 = nodes.pop()
            nodes.append(HuffmanTreeNode(None, n1.freq + n2.freq, n1, n2))
        
        self._coding_tree = nodes[0]
        self._generate_codebook(self._coding_tree, [])

    def _generate_codebook(self, node, codebook):
        if node.is_leaf():
            self.codebook[node.item] = "".join(codebook)
        else:
            if node.lchild is not None:
                self._generate_codebook(node.lchild, codebook + ['0'])
            
            if node.rchild is not None:
                self._generate_codebook(node.rchild, codebook + ['1'])

    def get_codebook(self):
        return str(self.codebook)

    def encode(self):
        output = []
        for ch in self.text:
            output.append(self.codebook[ch])
        
        return "".join(output)

    def decode(self, inputs):
        output = []
        if inputs == '':
            return self._coding_tree.item

        cur = self._coding_tree
        for c in inputs:
            cur = cur.lchild if c == '0' else cur.rchild
            if cur.is_leaf():
                output.append(cur.item)
                cur = self._coding_tree
                            
        return "".join(output)

if __name__ == "__main__":
    t = HuffmanTree()
    t.update("hello world")
    print(t.get_codebook())
    s = t.encode()
    print(s)
    print(t.decode(s))
    t.update('h')
    print(t.get_codebook())
    s = t.encode()
    print(s)
    print(t.decode(s))




    